{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please run\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "bin /usr/local/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so\n",
      "CUDA_SETUP: WARNING! libcudart.so not found in any environmental path. Searching in backup paths...\n",
      "CUDA SETUP: CUDA runtime path found: /usr/local/cuda/lib64/libcudart.so\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 117\n",
      "CUDA SETUP: Loading binary /usr/local/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/usr/local/nvidia/lib64'), PosixPath('/usr/local/nvidia/lib')}\n",
      "  warn(msg)\n",
      "/usr/local/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: /usr/local/nvidia/lib:/usr/local/nvidia/lib64 did not contain ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] as expected! Searching further paths...\n",
      "  warn(msg)\n",
      "/usr/local/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('//matplotlib_inline.backend_inline'), PosixPath('module')}\n",
      "  warn(msg)\n",
      "/usr/local/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/usr/local/cuda/lib64/libcudart.so'), PosixPath('/usr/local/cuda/lib64/libcudart.so.11.0')}.. We'll flip a coin and try one of these, in order to fail forward.\n",
      "Either way, this might cause trouble in the future:\n",
      "If you get `CUDA error: invalid device function` errors, the above might be the cause and the solution is to make sure only one ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] in the paths that we search based on your env.\n",
      "  warn(msg)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import LlamaForCausalLM, LlamaTokenizer\n",
    "\n",
    "model_id = \"../jinndata/数据代码/huggingface.co/seeledu/Chinese-Llama-2-7B\"\n",
    "\n",
    "tokenizer = LlamaTokenizer.from_pretrained(model_id)\n",
    "\n",
    "model = LlamaForCausalLM.from_pretrained(\n",
    "    model_id, load_in_8bit=True, device_map=\"auto\", torch_dtype=torch.float16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from configs.datasets import shennong_dataset\n",
    "from utils.dataset_utils import get_preprocessed_dataset\n",
    "\n",
    "train_dataset = get_preprocessed_dataset(tokenizer, shennong_dataset, \"train\")\n",
    "eval_dataset = get_preprocessed_dataset(tokenizer, shennong_dataset, \"val\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "你是一个中医临床医生，用正确的治疗方法来治疗疾病。\n",
      "你能够推荐常规药物、草药和其他天然替代品在提供建议时，还需要考虑患者的年龄、生活方式和病史。\n",
      "\n",
      "我的第一个建议请求是“患者出现半身不遂症状，没有其他症状。请推荐中药。”。\n",
      "\n",
      "建议: \n",
      "1. 考虑患者的年龄、生活方式和病史，然后推荐一种中药。\n"
     ]
    }
   ],
   "source": [
    "eval_prompt = \"\"\"\n",
    "你是一个中医临床医生，用正确的治疗方法来治疗疾病。\n",
    "你能够推荐常规药物、草药和其他天然替代品在提供建议时，还需要考虑患者的年龄、生活方式和病史。\\n\n",
    "我的第一个建议请求是“患者出现半身不遂症状，没有其他症状。请推荐中药。”。\\n\n",
    "建议: \n",
    "\"\"\"\n",
    "\n",
    "model_input = tokenizer(eval_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    raw_output = model.generate(**model_input, max_new_tokens=100)\n",
    "    print(tokenizer.decode(raw_output[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.9/site-packages/peft/utils/other.py:122: FutureWarning: prepare_model_for_int8_training is deprecated and will be removed in a future version. Use prepare_model_for_kbit_training instead.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 4,194,304 || all params: 6,742,618,112 || trainable%: 0.06220586618327525\n"
     ]
    }
   ],
   "source": [
    "model.train()\n",
    "\n",
    "\n",
    "def create_peft_config(model):\n",
    "    from peft import (\n",
    "        LoraConfig,\n",
    "        TaskType,\n",
    "        get_peft_model,\n",
    "        prepare_model_for_int8_training,\n",
    "    )\n",
    "\n",
    "    peft_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False,\n",
    "        r=8,\n",
    "        lora_alpha=32,\n",
    "        lora_dropout=0.05,\n",
    "        target_modules=[\"q_proj\", \"v_proj\"],\n",
    "    )\n",
    "\n",
    "    # prepare int-8 model for training\n",
    "    model = prepare_model_for_int8_training(model)\n",
    "    model = get_peft_model(model, peft_config)\n",
    "    model.print_trainable_parameters()\n",
    "    return model, peft_config\n",
    "\n",
    "\n",
    "# create peft config\n",
    "model, lora_config = create_peft_config(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from contextlib import nullcontext\n",
    "\n",
    "from transformers import TrainerCallback\n",
    "\n",
    "enable_profiler = False\n",
    "output_dir = \"outputs/llama-output\"\n",
    "\n",
    "config = {\n",
    "    \"lora_config\": lora_config,\n",
    "    \"learning_rate\": 1e-4,\n",
    "    \"num_train_epochs\": 1,\n",
    "    \"per_device_train_batch_size\": 4,\n",
    "    \"gradient_checkpointing\": False,\n",
    "}\n",
    "\n",
    "# Set up profiler\n",
    "if enable_profiler:\n",
    "    wait, warmup, active, repeat = 1, 1, 2, 1\n",
    "    total_steps = (wait + warmup + active) * (1 + repeat)\n",
    "    schedule = torch.profiler.schedule(\n",
    "        wait=wait, warmup=warmup, active=active, repeat=repeat\n",
    "    )\n",
    "    profiler = torch.profiler.profile(\n",
    "        schedule=schedule,\n",
    "        on_trace_ready=torch.profiler.tensorboard_trace_handler(\n",
    "            f\"{output_dir}/logs/tensorboard\"\n",
    "        ),\n",
    "        record_shapes=True,\n",
    "        profile_memory=True,\n",
    "        with_stack=True,\n",
    "    )\n",
    "\n",
    "    class ProfilerCallback(TrainerCallback):\n",
    "        def __init__(self, profiler):\n",
    "            self.profiler = profiler\n",
    "\n",
    "        def on_step_end(self, *args, **kwargs):\n",
    "            self.profiler.step()\n",
    "\n",
    "    profiler_callback = ProfilerCallback(profiler)\n",
    "else:\n",
    "    profiler = nullcontext()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'nltk'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[6], line 5\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfunctools\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m partial\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Trainer, TrainingArguments, default_data_collator\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmetric_utils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m similarity_metrics\n\u001b[1;32m      7\u001b[0m compute_metrics \u001b[38;5;241m=\u001b[39m partial(similarity_metrics, tokenizer\u001b[38;5;241m=\u001b[39mtokenizer)\n\u001b[1;32m      9\u001b[0m \u001b[38;5;66;03m# Set logger\u001b[39;00m\n",
      "File \u001b[0;32m~/zhongjing-mlm/utils/metric_utils.py:2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mnltk\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtranslate\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbleu_score\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m sentence_bleu, SmoothingFunction\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mrouge\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Rouge\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AutoTokenizer, EvalPrediction\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'nltk'"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from functools import partial\n",
    "\n",
    "from transformers import Trainer, TrainingArguments, default_data_collator\n",
    "from utils.metric_utils import similarity_metrics\n",
    "\n",
    "compute_metrics = partial(similarity_metrics, tokenizer=tokenizer)\n",
    "\n",
    "# Set logger\n",
    "os.eviron[\"WANDB_PROJECT\"] = \"ZhongJing\"  # log to your project\n",
    "# os.eviron[\"WANDB_LOG_MODEL\"] = \"all\" # log your models\n",
    "\n",
    "output_dir = \"../outputs\"\n",
    "# Define training args\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    overwrite_output_dir=True,\n",
    "    do_eval=True,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    per_device_eval_batch_size=1,\n",
    "    eval_steps=100,\n",
    "    bf16=True,  # Use BF16 if available\n",
    "    # logging strategies\n",
    "    logging_dir=f\"{output_dir}/logs\",\n",
    "    logging_strategy=\"steps\",\n",
    "    logging_steps=10,\n",
    "    save_strategy=\"no\",\n",
    "    optim=\"adamw_torch_fused\",\n",
    "    max_steps=total_steps if enable_profiler else -1,\n",
    "    report_to=\"wandb\",\n",
    "    **{k: v for k, v in config.items() if k != \"lora_config\"},\n",
    ")\n",
    "\n",
    "with profiler:\n",
    "    # Create Trainer instance\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        train_dataset=train_dataset,\n",
    "        eval_dataset=eval_dataset,\n",
    "        data_collator=default_data_collator,\n",
    "        callbacks=[profiler_callback] if enable_profiler else [],\n",
    "    )\n",
    "\n",
    "    # Start training\n",
    "    trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    print(\n",
    "        tokenizer.decode(\n",
    "            model.generate(**model_input, max_new_tokens=100)[0],\n",
    "            skip_special_tokens=True,\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ceval import LlamaEvaluator\n",
    "\n",
    "evaluator = LlamaEvaluator(model, tokenizer, fewshot_num=0)\n",
    "evaluator.evaluate_subjects(\"./evaluate/tasks\", output_dir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
